import dlib  # 人脸识别的库dlib
import cv2  # 图像处理的库OpenCv
import pandas as pd  # 数据处理的库Pandas
import os
import json

import numpy as np
from utils import *


def calc_face_feature(img_path):

    # face recognition model, the object maps human faces into 128D vectors
    facerec = dlib.face_recognition_model_v1("model/dlib_face_recognition_resnet_model_v1.dat")

    # Dlib 预测器
    detector = dlib.get_frontal_face_detector()
    predictor = dlib.shape_predictor('model/shape_predictor_68_face_landmarks.dat')


    files = os.listdir(img_path)
    feature_list = []
    feature_average = []
    is_success = False
    for file in files:
        fp = os.path.join(img_path, file)

        img = cv2_imread(fp, 0)
        img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        dets = detector(img_gray, 1)
        if len(dets) != 0:
            shape = predictor(img_gray, dets[0])
            face_descriptor = facerec.compute_face_descriptor(img_gray, shape)
            feature_list.append(face_descriptor)
        else:
            face_descriptor = 0
            # print("未在照片中识别到人脸")
    
    if len(feature_list) >= 5:
        for j in range(128):
            feature_average.append(0)
            for i in range(len(feature_list)):
                feature_average[j] += feature_list[i][j]
            feature_average[j] = (feature_average[j]) / len(feature_list)
        is_success = True
        feature_average = json.dumps(feature_average)
    else:
        is_success = False
        feature_average = None

    return is_success, feature_average


if __name__ == '__main__':
    fold = r"C:\Users\admin\Documents\pie-face-recognition-system\backend\app\dataset\train_data\s1"
    is_success, feature_average = calc_face_feature(fold)
    print(type(json.loads(feature_average)))